# Copyright 2020 Makani Technologies LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utilities for working with shell processes.

This module provides interfaces for launching subprocesses that would otherwise
be invoked via the shell.  It is intended to reduce reliance on Bash scripts and
to facilitate dependency injection in tests.
"""
# TODO: There are some nice opportunities for context managers here,
# particularly with Executor.RunInBackground() and Packager.OpenPackage().
# Figure out if that can be done without making mocks too burdensome.

import io
import logging
import os
import signal
import subprocess
import tarfile
import time

from makani.lib.python import os_util


class ExecutorError(Exception):
  pass


class TimeoutError(Exception):
  pass


# NOTE: Prefer direct use of the subprocess module to this class where
# possible. I wrote it to facilitate dependency injection before I understood
# the mock module.
class Executor(object):
  """Interface to processes that invoke a specific executable.

  Processes may be started either in the foreground, via Run(), or in the
  background, via RunInBackground().

  When the Executor is destroyed, any outstanding background processes will be
  terminated.
  """

  def __init__(self):
    self._background_processes = set()

  def __del__(self):
    # self._background_processes is modified by KillIfRunning(), so we have to
    # copy its contents first.
    for process in list(self._background_processes):
      self.KillIfRunning(process)

  def Run(self, popen_args, timeout_sec=None, **kwargs):
    """Runs the executable in the foreground with the provided args.

    Args:
      popen_args: List of Popen-style args.
      timeout_sec: Number of seconds after which to kill the process. Polling
          is performed at 5 second intervals.
      **kwargs: Forwarded to Popen.

    Returns:
      If stdout=subprocess.PIPE or stderr=subprocess.PIPE is specified in
      popen_kwargs, then this method returns:
          (<subprocess.Popen>, <str>, <str>): Completed process, stdout, stderr.
      Otherwise:
          <subprocess.Popen>: Completed process.

    Raises:
      TimeoutError: The subprocess exceeds a time given by timeout_sec.
    """

    popen = subprocess.Popen(popen_args, **kwargs)

    return_streams = (kwargs.get('stdout', None) == subprocess.PIPE or
                      kwargs.get('stderr', None) == subprocess.PIPE)

    if timeout_sec is not None:
      end_time = time.time() + timeout_sec
      while popen.poll() is None:
        if time.time() >= end_time:
          popen.send_signal(signal.SIGINT)
          raise TimeoutError(
              'Command line "%s" exceeded time limit of %d seconds.'
              % (popen_args, timeout_sec))
        time.sleep(5)

    stdout, stderr = popen.communicate()
    if return_streams:
      return popen, stdout, stderr
    else:
      return popen

  def CheckRun(self, popen_args, **kwargs):
    """Runs the executable in the foreground and checks successful completion.

    Args:
      popen_args: List of Popen-style args.
      **kwargs: Forwarded to Popen.

    Returns:
      See Run().

    Raises:
      ExecutorError: The process exited with nonzero return code.
    """

    output = self.Run(popen_args, **kwargs)

    if isinstance(output, subprocess.Popen):
      if output.returncode != 0:
        raise ExecutorError(
            'Process generated by %s exited with return code %d.'
            % (popen_args, output.returncode))
    else:
      popen, _, stderr = output  # pylint: disable=unpacking-non-sequence
      if popen.returncode != 0:
        raise ExecutorError(
            'Process generated by %s exited with return code %d.  stderr:\n%s'
            % (popen_args, popen.returncode, stderr))

    return output

  def RunInBackground(self, popen_args, **kwargs):
    """Runs the executable in the background with the provided args.

    The background process is owned by this Executor until KillIfRunning() is
    called on it.

    stdout and stderr are directed to subprocess.PIPE, so they will be available
    via Popen.communicate() when the process completes.

    Args:
      popen_args: List of Popen-style args.
      **kwargs: Forwarded to Popen.

    Returns:
       <subprocess.Popen> The background process.
    """

    process = subprocess.Popen(popen_args, **kwargs)
    self._background_processes.add(process)
    return process

  def KillIfRunning(self, process):
    """Kills a background process with SIGINT if it is still running.

    Also removes the process from the Executor's background processes.

    Args:
      process: A background process currently owned by the Executor.

    Raises:
      ExecutorError: The process is not owned by the Executor.
    """

    if process not in self._background_processes:
      raise ExecutorError('This Executor does not own process %d.' %
                          process.pid)
    self._background_processes.remove(process)

    if process.poll() is None:
      try:
        process.send_signal(signal.SIGINT)
        process.communicate()
      except OSError as e:
        if e.errno == os.errno.ESRCH:
          logging.warning('Process %d was terminated while killing.',
                          process.pid)
        else:
          raise


class PackagerException(Exception):
  pass


class Packager(object):
  """Creates gzipped tar archives, a.k.a. "packages"."""

  def __init__(self):
    self._tarball = None

  def CreateRepoPackage(self, package_path,
                        commit='HEAD',
                        restrict_paths=None):
    """Creates a package from HEAD of the git respository.

    Args:
      package_path: Path of the package file.
      commit: Commit for which to produce the archive.
      restrict_paths: List of paths, relative to makani.HOME, to which the
          archive will be restricted.
    """
    if restrict_paths is None:
      restrict_paths = []
    executor = Executor()

    # If running under Bazel, we're working with a symlink to this file. In that
    # case, we get the absolute path to this file in the repo by changing .pyc
    # to .py if necessary, then resolving any symlinks. The we use "git
    # rev-parse --show-toplevel" from the containing directory to determine the
    # base directory of the repo.
    abs_path = os.path.realpath(__file__[:-1] if __file__.endswith('.pyc')
                                else __file__)
    with os_util.ChangeDir(os.path.dirname(abs_path)):
      repo_base = executor.CheckRun(['git', 'rev-parse', '--show-toplevel'],
                                    stdout=subprocess.PIPE)[1].strip()

    with os_util.ChangeDir(repo_base):
      # Warn if there are uncommitted changes.
      status = executor.CheckRun(['git', 'status'], stdout=subprocess.PIPE)[1]
      if commit == 'HEAD' and 'Changes not staged for commit' in status:
        logging.warning('You have uncommited changes.  Only committed changes '
                        'will be included in the repository package.')

      executor.CheckRun(['git', 'archive', '--format=tar.gz',
                         '--output=' + package_path, commit] + restrict_paths)

  def OpenPackage(self, package_path):
    """Starts a new package.

    Args:
      package_path: Path of the output package.

    Raises:
      PackagerException: A package is already open.
    """
    if self._tarball is not None:
      raise PackagerException('Must close current package with ClosePackage '
                              'before creating a new one.')
    self._tarball = tarfile.open(package_path, 'w:gz')

  def _CheckOpened(self):
    if self._tarball is None:
      raise PackagerException('No package has been opened with OpenPackage().')

  def AddFile(self, file_name, name_in_archive):
    """Adds a file, specified by name, to the package.

    Args:
      file_name: Name of the file to add.
      name_in_archive: Name of the new object in the archive.

    Raises:
      PackagerException: No package is open.
    """
    self._CheckOpened()
    self._tarball.add(file_name, arcname=name_in_archive)

  def AddString(self, contents, name_in_archive, mode=0644):
    """Adds the contents of a string to a new file in the archive.

    Args:
      contents: String of contents for the new file.
      name_in_archive: Name of the new file in the archive.
      mode: Permissions for the file.

    Raises:
      PackagerException: No package is open.
    """
    self._CheckOpened()

    tarinfo = tarfile.TarInfo(name_in_archive)
    tarinfo.size = len(contents)
    tarinfo.mtime = time.time()
    tarinfo.mode = mode
    self._tarball.addfile(tarinfo, io.BytesIO(contents))

  def ClosePackage(self):
    """Completes the current package, after which it may be used externally.

    Raises:
      PackagerException: No package is open.
    """
    self._CheckOpened()

    self._tarball.close()
    self._tarball = None


def GetStdout(cmd):
  """Gets stdout for a command.

  Args:
    cmd: Shell command, as a string or list of strings. It should return 0.

  Returns:
    stdout from the command.

  Raises:
    RuntimeError: If the return code was not 0.
  """
  if isinstance(cmd, str):
    cmd = cmd.split()
  popen = subprocess.Popen(cmd, stdout=subprocess.PIPE)
  stdout, stderr = popen.communicate()
  if popen.returncode:
    raise RuntimeError('Command "%s" failed with return code %d. stderr:\n%s'
                       % (cmd, popen.returncode, stderr))
  return stdout
